#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import json
import argparse
from typing import Dict, List

class ExperimentResultsAnalyzer:
    """
    实验结果分析器，用于查看和比较不同窗口大小实验组的平均指标
    """
    
    def __init__(self, results_dir: str):
        """
        初始化分析器
        
        Args:
            results_dir: 实验结果文件夹路径
        """
        self.results_dir = results_dir
    
    def load_experiment_results(self) -> Dict[str, Dict]:
        """
        加载所有实验结果文件
        
        Returns:
            窗口大小到实验结果的映射字典
        """
        experiment_results = {}
        
        # 遍历结果文件夹中的所有文件
        for filename in os.listdir(self.results_dir):
            if filename.startswith('experiment_results_window_') and filename.endswith('.json'):
                # 提取窗口大小
                window_size = filename.replace('experiment_results_window_', '').replace('.json', '')
                
                # 尝试将窗口大小转换为数字（如果是数字的话）
                try:
                    window_size = int(window_size)
                except ValueError:
                    pass  # 保持原字符串
                
                file_path = os.path.join(self.results_dir, filename)
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                        if data and isinstance(data, list) and len(data) > 0:
                            # 获取第一个实验组的数据（应该只有一个）
                            experiment_results[window_size] = data[0]
                            print(f"已加载窗口大小 {window_size} 的实验结果")
                except Exception as e:
                    print(f"加载文件 {filename} 失败: {str(e)}")
        
        return experiment_results
    
    def print_average_metrics_comparison(self, experiment_results: Dict[str, Dict]):
        """
        以表格形式打印不同实验组的平均指标比较
        
        Args:
            experiment_results: 窗口大小到实验结果的映射字典
        """
        if not experiment_results:
            print("没有找到实验结果文件")
            return
        
        # 获取所有窗口大小并排序，处理混合类型
        window_sizes = list(experiment_results.keys())
        # 先尝试将所有可以转换为整数的转换为整数，以便正确排序
        sorted_window_sizes = []
        string_windows = []
        
        for window in window_sizes:
            try:
                sorted_window_sizes.append(int(window))
            except (ValueError, TypeError):
                string_windows.append(str(window))
        
        # 对数字窗口和字符串窗口分别排序，然后合并
        sorted_window_sizes.sort()
        string_windows.sort()
        
        window_sizes = sorted_window_sizes + string_windows
        
        # 收集所有可能的指标名称
        all_metrics = set()
        for window_size, result in experiment_results.items():
            if 'average_metrics' in result:
                all_metrics.update(result['average_metrics'].keys())
        
        # 排序指标名称，把常见的评估指标放在前面
        sorted_metrics = sorted(all_metrics)
        
        # 打印表头
        print("\n各窗口大小实验组平均指标比较：")
        header = "窗口大小\t"
        for metric in sorted_metrics:
            header += f"{metric}\t"
        print(header)
        
        # 打印分隔线
        separator = "-------\t"
        for _ in sorted_metrics:
            separator += "-------\t"
        print(separator)
        
        # 打印每个窗口大小的指标值
        for window_size in window_sizes:
            result = experiment_results[window_size]
            row = f"{window_size}\t"
            
            if 'average_metrics' in result:
                metrics = result['average_metrics']
                for metric in sorted_metrics:
                    if metric in metrics:
                        row += f"{metrics[metric]:.4f}\t"
                    else:
                        row += "N/A\t"
            else:
                for _ in sorted_metrics:
                    row += "N/A\t"
            
            print(row)
    
    def export_summary_to_file(self, experiment_results: Dict[str, Dict], output_file: str):
        """
        将平均指标比较结果导出到文件
        
        Args:
            experiment_results: 窗口大小到实验结果的映射字典
            output_file: 输出文件路径
        """
        summary = {}
        for window_size, result in experiment_results.items():
            if 'average_metrics' in result:
                summary[str(window_size)] = result['average_metrics']
            
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(summary, f, ensure_ascii=False, indent=2)
        
        print(f"\n平均指标比较结果已导出到: {output_file}")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='分析不同窗口大小实验组的平均指标')
    parser.add_argument('--results-dir', type=str, default='../results', 
                        help='实验结果文件夹路径')
    parser.add_argument('--export', type=str, default=None, 
                        help='是否将结果导出到文件，指定文件路径')
    
    args = parser.parse_args()
    
    # 创建分析器实例
    analyzer = ExperimentResultsAnalyzer(args.results_dir)
    
    # 加载实验结果
    experiment_results = analyzer.load_experiment_results()
    
    # 打印平均指标比较
    analyzer.print_average_metrics_comparison(experiment_results)
    
    # 如果指定了导出文件，则导出结果
    if args.export:
        analyzer.export_summary_to_file(experiment_results, args.export)


if __name__ == '__main__':
    main()